from hurryoff.db import *
from sqlalchemy import or_
from sqlalchemy.orm import contains_eager

class Router:

    def __init__(self, startPoint, destination, filters=None):
        self.routes = []
        self.passed = []
        self.filters = filters
        self.startPoint = startPoint
        self.destination = destination
        self.passed.append(self.startPoint)

    def getRoutes(self):
        self.__initialise()
        while len(self.routes) > 0:
            added = []
            for routesSet in self.routes:
                currentLocation = routesSet[len(routesSet) - 1].destinationPoint
                if currentLocation == self.destination:
                    yield self.__preparePath(routesSet)
                    self.routes.remove(routesSet)
                    continue
                newWorkingSet = [path for path in self.__getWorkingSetForLocation(currentLocation) if not self.__isVisited(path)]
                if len(newWorkingSet) > 0:
                    for part in newWorkingSet[1:]:
                        newRoute = routesSet[:]
                        newRoute.append(part)
                        added.append(newRoute)
                    routesSet.append(newWorkingSet[0])
                    self.passed.append(newWorkingSet[0].startPoint)
                else:
                    self.routes.remove(routesSet)
            self.routes.extend(added)
            if self.filters:
                for filter in self.filters:
                    self.routes = [route for route in self.routes if filter(self.__preparePath(route))]

    def __isVisited(self, location):
        if location is None:
            return True
        return location.destinationPoint in self.passed

    def __preparePath(self, parts):
        result = [parts[0]]
        for part in parts[1:len(parts)]:
            if part.trainId != result[len(result) - 1].trainId:
                result.append(part)
            else:
                result[len(result) - 1].destinationPoint = part.destinationPoint
                result[len(result) - 1].destinationTime = part.destinationTime
        return result

    def __initialise(self):
        for path in [path for path in self.__getWorkingSetForLocation(self.startPoint) if path is not None]:
            self.routes.append([path, ])

    def __getWorkingSetForLocation(self, location):
        workingSet = []
        isDirectTrainPresent, trains = isDirectlyConnected(location, self.destination)
        if isDirectTrainPresent:
            for trainId in trains:
                workingSet.append(getPathPart(location, self.destination, trainId))
        for (trainId, station) in getNeighboring(location).items():
            workingSet.append(getPathPart(location, station.name, trainId))
        return workingSet

class PathPart:

    def __init__(self, startPoint, destinationPoint, departureTime, destinationTime, trainId):
        self.startPoint = startPoint
        self.destinationPoint = destinationPoint
        self.departureTime = departureTime
        self.destinationTime = destinationTime
        self.trainId = trainId

    def __eq__(self, other):
        if other is None:
            return False
        return (self.startPoint == other.startPoint) and (self.destinationPoint == other.destinationPoint) and \
               (self.departureTime == other.departureTime) and (self.destinationTime == other.destinationTime) and \
               (self.trainId == other.trainId)

def isDirectlyConnected(location1, location2):
    session = Session()
    try:
        q1 = session.query(Train.id)\
                        .join(Train.stops, Stop.station) \
                        .filter(Station.name == location1) \
                        .subquery()
        q2 = session.query(Train.id)\
                        .join(Train.stops, Stop.station)\
                        .filter(Station.name == location2)
        directTrains = q2.filter(Train.id.in_(q1)).all()
        return len(directTrains) > 0, [train[0] for train in directTrains]
    finally:
        session.close()

def getNeighboring(locationName):
    neighbors = {}
    session = Session()
    try:
        q1 = session.query(Train.id, Stop.order)\
                        .join(Train.stops, Stop.station) \
                        .filter(Station.name == locationName) \
                        .all()
        for (train_id, stop_order) in q1:
            station = session.query(Station)\
                                .join(Train.stops, Stop.station) \
                                .filter(Stop.order == stop_order + 1)\
                                .filter(Train.id == str(train_id))\
                                .first()
            if station:
                neighbors[train_id] = station
        return neighbors
    finally:
        session.close()

def getPathPart(departurePoint, destinationPoint, trainId):
    session = Session()
    try:
        stops = session.query(Stop)\
                    .join(Train.stops, Stop.station)\
                    .options(contains_eager(Stop.station))\
                    .filter(or_(Station.name == departurePoint, Station.name == destinationPoint))\
                    .filter(Train.id == str(trainId))\
                    .order_by(Stop.order)\
                    .all()
        if stops[0].station.name == departurePoint:
            return PathPart(departurePoint, destinationPoint, stops[0].departure, stops[1].arrival, trainId)
        return None
    finally:
        session.close()

